#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
from PyQt5.QtWidgets import QMessageBox, QDialog
import platform
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure

from utils.math import calculate_displacement, smooth_and_recalculate
from view.custom_dialog import PlotDialog


class ThreeDPlotWidget(FigureCanvas):
    def __init__(self, origin_data, draw_widget, width=12, height=100, dpi=100):
        # 根据操作系统类型设置不同的字体方案
        if platform.system() == 'Darwin':  # macOS系统
            plt.rcParams['font.family'] = ['Hei', 'Arial', 'Helvetica', 'Times New Roman']
        else:  # 其他系统
            plt.rcParams['font.family'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)

        # 创建 3D 坐标轴
        self.ax = self.fig.add_subplot(111, projection='3d')

        self.origin_data = origin_data
        self.Draw_widget = draw_widget
        # 存储初始视角
        self.initial_elev = 20
        self.initial_azim = 45
        self.ax.view_init(elev=self.initial_elev, azim=self.initial_azim)

        # 添加 colorbar 引用
        self.cbar = None

        # 存储数据范围用于缩放
        self.data_limits = {
            'xlim': None,
            'ylim': None,
            'zlim': None
        }
        # 平滑还是不平滑
        self.current_plot_type = 1

        # 当前缩放级别
        self.zoom_level = 1.0
        self.zoom_factor = 1.2

        # 悬浮提示相关
        self.annotation = None
        self.current_df = None
        self.scatter_points = None

        # 多标注点相关属性
        self.marked_points = []  # 存储所有标注点信息，每个元素为字典
        self.marked_point_markers = []  # 存储所有标注点的标记对象
        self.marked_points_annotation = None  # 左下角固定标注点信息

        # 连接鼠标移动事件
        self.mpl_connect('pick_event', self.on_pick)

    def safe_clear(self):
        """安全清除图形内容"""
        # 清除之前的colorbar
        if self.cbar is not None:
            try:
                self.cbar.remove()
            except:
                pass
            self.cbar = None

        # 清除坐标轴内容
        if self.ax is not None:
            self.ax.clear()

        # 清除整个图形
        self.fig.clear()

        # 重新创建坐标轴
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.ax.view_init(elev=self.initial_elev, azim=self.initial_azim)

        # 重置缩放级别
        self.zoom_level = 1.0

        # 重置标注点标记
        self.marked_point_markers = []

        # 重新连接事件
        self.mpl_connect('pick_event', self.on_pick)

        # 绘制一个空白的图形
        self.ax.set_xlabel('东移 (m)')
        self.ax.set_ylabel('北移 (m)')
        self.ax.set_zlabel('垂深 (m)')
        self.ax.set_title('井眼轨迹 3D 图')

        # 更新显示
        self.draw()

    def plot_well_trajectory_with_dogleg(self, df, use_dogleg_severity=True):
        """绘制井眼轨迹 3D 图，可选择使用狗腿度作为颜色映射"""
        # 清除之前的图形和colorbar
        self.safe_clear()

        # 保存数据用于悬浮提示
        self.current_df = df.copy()

        # 提取数据
        north = df['北移'].values
        east = df['东移'].values
        depth = df['井垂深'].values

        if use_dogleg_severity:
            # 使用狗腿度作为颜色映射
            colors = df['狗腿度'].values
            cmap = 'hot'
            color_label = '狗腿度 (°/30m)'
            title_suffix = ' - 颜色表示狗腿度'
        else:
            # 使用井斜深作为颜色映射（原始方式）
            colors = df['井斜深'].values
            cmap = 'viridis'
            color_label = '井斜深 (m)'
            title_suffix = ''

        # 绘制 3D 轨迹线
        self.scatter_points = self.ax.scatter(east, north, -depth,
                                              c=colors, cmap=cmap,
                                              s=30, alpha=0.8, picker=True)

        # 绘制连线
        self.ax.plot(east, north, -depth, 'k-', alpha=0.5, linewidth=1)

        # 标记高狗腿度区域（如果使用狗腿度）
        if use_dogleg_severity and len(df) > 0:
            high_dls_threshold = np.percentile(df['狗腿度'].values, 90)
            high_dls_mask = df['狗腿度'] > high_dls_threshold

            if np.any(high_dls_mask):
                self.ax.scatter(east[high_dls_mask], north[high_dls_mask], -depth[high_dls_mask],
                                color='red', s=50, alpha=0.9,
                                label=f'高狗腿度区域(>{high_dls_threshold:.1f}°/30m)')

        # 标记起点和终点
        self.ax.scatter(east[0], north[0], -depth[0],
                        color='green', s=100, label='起点', marker='o')
        self.ax.scatter(east[-1], north[-1], -depth[-1],
                        color='blue', s=100, label='终点', marker='s')

        # 标记所有标注点
        self._mark_all_points(df)

        # 设置坐标轴标签
        self.ax.set_xlabel('东移 (m)')
        self.ax.set_ylabel('北移 (m)')
        self.ax.set_zlabel('垂深 (m)')

        # 设置标题
        title = f'井眼轨迹 3D 图{title_suffix}'

        self.ax.set_title(title)

        # 添加颜色条
        self.cbar = self.fig.colorbar(self.scatter_points, ax=self.ax, shrink=0.6, aspect=20)
        self.cbar.set_label(color_label)

        # 添加图例
        self.ax.legend()

        # 存储数据范围用于缩放
        self._store_data_limits(east, north, depth)

        # 应用当前缩放级别
        self._apply_zoom()

        # 重置为初始视角
        self.ax.view_init(elev=self.initial_elev, azim=self.initial_azim)

        # 在左下角显示所有标注点信息
        self._show_marked_points_info()

        self.draw()

    def _mark_all_points(self, df):
        """标注所有标记点"""
        self.marked_point_markers = []

        for i, point in enumerate(self.marked_points):
            # 找到最接近标注点深度的数据点
            depth_diff = np.abs(df['井斜深'] - point['depth'])
            closest_idx = depth_diff.idxmin()

            if depth_diff.iloc[closest_idx] < 100.0:  # 如果深度差小于100米，认为是有效点
                point_data = df.iloc[closest_idx]

                # 使用不同的颜色和标记来区分不同的标注点
                colors = ['red', 'blue', 'orange', 'purple', 'brown']
                markers = ['*', 'D', 's', '^', 'v']

                color = colors[i % len(colors)]
                marker = markers[i % len(markers)]

                # 绘制标注点标记
                marker_obj = self.ax.scatter(
                    point_data['东移'],
                    point_data['北移'],
                    -point_data['井垂深'],
                    color=color,
                    s=200,
                    label=f"{point['description']}",
                    marker=marker,
                    edgecolors='black',
                    linewidth=2,
                    zorder=10  # 确保在最上层显示
                )
                self.marked_point_markers.append(marker_obj)

    def _show_marked_points_info(self):
        """在左下角显示所有标注点信息"""
        # 移除旧的标注点信息
        if self.marked_points_annotation:
            try:
                self.marked_points_annotation.remove()
            except:
                pass

        if not self.marked_points:
            return

        # 创建标注点信息文本
        info_text = "标注点信息:\n"
        for i, point in enumerate(self.marked_points):
            info_text += f"点{i + 1}: 深度{point['depth']}m\n"
            if 'description' in point:
                info_text += f"  描述: {point['description']}\n"

        # 在图形左下角显示标注点信息
        self.marked_points_annotation = self.fig.text(
            0.02, 0.3, info_text,
            transform=self.fig.transFigure,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue",
                      edgecolor="blue", alpha=0.9),
            fontsize=8,
            verticalalignment='top',
            zorder=1000
        )

    def add_marked_point(self, depth, description=None):
        """添加一个标注点"""
        point_info = {
            'depth': depth,
            'description': description
        }
        self.marked_points.append(point_info)

        # 如果已经有数据，重新绘制图形
        if self.current_df is not None:
            self.plot_well_trajectory_with_dogleg(self.current_df, True)

    def set_marked_points(self, points_list):
        """设置多个标注点，points_list为深度列表或包含深度和描述的元组列表"""
        self.marked_points = []

        for point in points_list:
            if isinstance(point, (int, float)):
                # 如果只提供深度
                self.add_marked_point(point)
            elif isinstance(point, (tuple, list)) and len(point) >= 1:
                # 如果提供深度和描述
                depth = point[0]
                description = point[1] if len(point) > 1 else None
                self.add_marked_point(depth, description)

    def clear_marked_points(self):
        """清除所有标注点标记"""
        self.marked_points = []
        self.marked_point_markers = []

        # 重新绘制
        if self.current_df is not None:
            self.plot_well_trajectory_with_dogleg(self.current_df, True)

    def remove_marked_point(self, index):
        """移除指定索引的标注点"""
        if 0 <= index < len(self.marked_points):
            self.marked_points.pop(index)
            # 重新绘制
            if self.current_df is not None:
                self.plot_well_trajectory_with_dogleg(self.current_df, True)

    def _store_data_limits(self, east, north, depth):
        """存储数据的原始范围"""
        margin = 0.1  # 10% 的边距

        x_range = np.ptp(east)
        y_range = np.ptp(north)
        z_range = np.ptp(depth)

        self.data_limits = {
            'xlim': (east.min() - margin * x_range, east.max() + margin * x_range),
            'ylim': (north.min() - margin * y_range, north.max() + margin * y_range),
            'zlim': (-depth.max() - margin * z_range, -depth.min() + margin * z_range)
        }

    def _apply_zoom(self):
        """应用当前的缩放级别"""
        if not all(limit is not None for limit in self.data_limits.values()):
            return

        x_center = np.mean(self.data_limits['xlim'])
        y_center = np.mean(self.data_limits['ylim'])
        z_center = np.mean(self.data_limits['zlim'])

        x_range = (self.data_limits['xlim'][1] - self.data_limits['xlim'][0]) / self.zoom_level
        y_range = (self.data_limits['ylim'][1] - self.data_limits['ylim'][0]) / self.zoom_level
        z_range = (self.data_limits['zlim'][1] - self.data_limits['zlim'][0]) / self.zoom_level

        self.ax.set_xlim(x_center - x_range / 2, x_center + x_range / 2)
        self.ax.set_ylim(y_center - y_range / 2, y_center + y_range / 2)
        self.ax.set_zlim(z_center - z_range / 2, z_center + z_range / 2)

    def on_pick(self, event):
        """处理pick事件"""
        if event.artist != self.scatter_points:
            return

        # 获取被选中的点的索引
        ind = event.ind[0]

        # 检查点击的是否是标注点附近
        clicked_depth = self.current_df.iloc[ind]['井斜深']
        is_marked_point = False
        marked_point_index = -1

        for i, point in enumerate(self.marked_points):
            depth_diff = abs(clicked_depth - point['depth'])
            if depth_diff < 1.0:  # 如果点击的点接近标注点
                is_marked_point = True
                marked_point_index = i
                break

        if is_marked_point:
            self.show_marked_point_annotation(marked_point_index)
        else:
            self.show_pick_annotation(ind)

        self.draw()

    def show_pick_annotation(self, index):
        """显示选中点的数据注释（在右下角）"""
        if self.current_df is None or index >= len(self.current_df):
            return

        # 移除旧注释
        if self.annotation:
            self.annotation.remove()

        # 获取数据
        row = self.current_df.iloc[index]

        # 创建注释文本
        text = (f"测点: {index}\n"
                f"井斜深: {row.get('井斜深', 'N/A'):.2f} m\n"
                f"井垂深: {row.get('井垂深', 'N/A'):.2f} m\n"
                f"北移: {row.get('北移', 'N/A'):.2f} m\n"
                f"东移: {row.get('东移', 'N/A'):.2f} m\n"
                f"井斜角: {row.get('井斜角', 'N/A'):.2f}°\n"
                f"方位角: {row.get('方位角', 'N/A'):.2f}°")

        if '狗腿度' in row:
            text += f"\n狗腿度: {row['狗腿度']:.2f}°/30m"

        # 在图形右下角显示注释
        self.annotation = self.fig.text(
            0.98, 0.02, text,
            transform=self.fig.transFigure,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow",
                      edgecolor="gray", alpha=0.9),
            fontsize=9,
            verticalalignment='bottom',
            horizontalalignment='right',
            zorder=1000
        )

    def show_marked_point_annotation(self, point_index):
        """显示标注点的特殊注释（在右下角）"""
        if point_index >= len(self.marked_points):
            return

        # 移除旧注释
        if self.annotation:
            self.annotation.remove()

        # 获取标注点信息
        point = self.marked_points[point_index]

        # 找到对应的数据点
        depth_diff = np.abs(self.current_df['井斜深'] - point['depth'])
        closest_idx = depth_diff.idxmin()
        row = self.current_df.iloc[closest_idx]

        # 创建标注点特殊注释文本
        text = (f"标注点 {point_index + 1}\n"
                f"井斜深: {row.get('井斜深', 'N/A'):.2f} m\n"
                f"井垂深: {row.get('井垂深', 'N/A'):.2f} m\n"
                f"北移: {row.get('北移', 'N/A'):.2f} m\n"
                f"东移: {row.get('东移', 'N/A'):.2f} m\n"
                f"井斜角: {row.get('井斜角', 'N/A'):.2f}°\n"
                f"方位角: {row.get('方位角', 'N/A'):.2f}°")

        if '狗腿度' in row:
            text += f"\n狗腿度: {row['狗腿度']:.2f}°/30m"

        if point.get('description'):
            text += f"\n描述: {point['description']}"

        # 在图形右下角显示注释
        self.annotation = self.fig.text(
            0.98, 0.02, text,
            transform=self.fig.transFigure,
            bbox=dict(boxstyle="round,pad=0.3", facecolor="lightcoral",
                      edgecolor="red", alpha=0.9),
            fontsize=9,
            verticalalignment='bottom',
            horizontalalignment='right',
            zorder=1000
        )

    def zoom_in(self):
        """放大视图"""
        if self.zoom_level < 10:
            self.zoom_level *= self.zoom_factor
            self._apply_zoom()
            self.draw()

    def zoom_out(self):
        """缩小视图"""
        if self.zoom_level > 0.1:
            self.zoom_level /= self.zoom_factor
            self._apply_zoom()
            self.draw()

    def reset_view(self):
        """重置视图到初始状态"""
        self.zoom_level = 1.0
        self._apply_zoom()
        self.ax.view_init(elev=self.initial_elev, azim=self.initial_azim)
        self.draw()

    def clear_plot(self):
        """清除当前绘制的图片"""
        # 清除绘图区域
        self.clear_draw_widget()
        # 重置当前绘图类型
        self.current_plot_type = None
        # 如果有3D绘图部件，清除其内容
        self.safe_clear()

    def show_plot_dialog(self, window):
        """显示绘图选择对话框"""
        dialog = PlotDialog(window)
        if dialog.exec_() == QDialog.Accepted:
            plot_type = dialog.get_selected_plot_type()
            self.plot_data(plot_type)

    def plot_data(self, plot_type):
        """根据选择的类型在主界面右侧绘制数据图表"""
        # 存储当前绘图类型
        self.current_plot_type = plot_type

        # 确保有数据
        if self.origin_data is None:
            QMessageBox.warning(None, "错误", "没有数据可供绘图")
            return
        # 清除之前的绘图
        self.clear_draw_widget()

        self.plot_3d_trajectory()

    def plot_3d_trajectory(self):
        """绘制3D井眼轨迹图"""
        # 计算视位移、北移和东移
        if self.current_plot_type == 1:
            df = calculate_displacement(self.origin_data)
        else:
            df = smooth_and_recalculate(self.origin_data)

        self.Draw_widget.layout().addWidget(self)
        # 显示3D绘图部件
        self.show()
        # 绘制3D轨迹
        self.plot_well_trajectory_with_dogleg(df, True)

    def clear_draw_widget(self):
        """清除绘图区域"""
        layout = self.Draw_widget.layout()
        if layout:
            # 隐藏所有部件
            for i in range(layout.count()):
                item = layout.itemAt(i)
                if item.widget():
                    item.widget().hide()